import math
import numpy as np
import scipy.linalg as la
import time
from scipy.spatial import cKDTree
import os

dir_current_file = os.path.dirname(__file__)  # 'Algorithm_demo_1\planner'
dir_parent_1 = os.path.dirname(dir_current_file)  # 'Algorithm_demo_1'

npz_path = os.path.join(dir_current_file, "all_results", "lqr_K", "lqr_K_table.npz")


# State 对象表示自车的状态，位置x、y，以及横摆角yaw、速度v
class State:

    def __init__(self, x=0.0, y=0.0, yaw=0.0, v=0.0):
        self.x = x
        self.y = y
        self.yaw = yaw
        self.v = v

def count_repeats(ind_list):
    if len(ind_list) < 2:
        return 1

    count = 1
    for i in range(len(ind_list) - 2, -1, -1):
        if ind_list[-1] == ind_list[i]:
            count += 1
        else:
            break

    return count

class LQRController:

    def __init__(self):
        self.dt = 0.1  # time tick[s]，采样时间
        self.L = 5.6  # Wheel base of the vehicle [m]，车辆轴距
        self.max_steer =1#np.deg2rad(90)  # maximum steering angle[rad]
        # LQR parameter
        self.lqr_Q = np.eye(5)
        # self.lqr_Q[0][0] = 2
        # self.lqr_Q[1][1] = 2
        # self.lqr_Q[2][2] = 2
        # self.lqr_Q[3][3] = 2
        # self.lqr_Q[4][4] = 5
        self.lqr_R = np.eye(2)
        self.pe, self.pth_e = 0.0, 0.0
        self.keyind = 0 
        self.isback = False
        data = np.load(npz_path)
        self.v_tbl = data["v_tbl"]    # shape = (N,)
        self.K_tbl = data["K_tbl"]  

    def pi_2_pi(self,angle):
        return (angle + math.pi) % (2 * math.pi) - math.pi

    # 实现离散Riccati equation 的求解方法
    def solve_dare(self,A, B, Q, R):
        """
        solve a discrete time_Algebraic Riccati equation (DARE)
        """
        x = Q
        x_next = Q
        max_iter = 150
        eps = 0.01

        for i in range(max_iter):
            x_next = A.T @ x @ A - A.T @ x @ B @ \
                    la.inv(R + B.T @ x @ B) @ B.T @ x @ A + Q
            if (abs(x_next - x)).max() < eps:
                break
            x = x_next

        return x_next

    # 返回值K 即为LQR 问题求解方法中系数K的解
    def dlqr(self,A, B, Q, R):
        """Solve the discrete time lqr controller.
        x[k+1] = A x[k] + B u[k]
        cost = sum x[k].T*Q*x[k] + u[k].T*R*u[k]
        # ref Bertsekas, p.151
        """

        # first, try to solve the ricatti equation
        X = self.solve_dare(A, B, Q, R)

        # compute the LQR gain
        K = la.inv(B.T @ X @ B + R) @ (B.T @ X @ A)

        eig_result = la.eig(A - B @ K)

        return K, X, eig_result[0]

    def get_K(self, v):
        # 1) 限定 v 在表格范围内
        v_min, v_max = self.v_tbl[0], self.v_tbl[-1]
        if v <= v_min:
            return self.K_tbl[0]
        if v >= v_max:
            return self.K_tbl[-1]
        # 2) 找到 v 在 v_tbl 中的插值区间
        idx = np.searchsorted(self.v_tbl, v)  # 第一个大于等于 v 的索引
        v_lo, v_hi = self.v_tbl[idx-1], self.v_tbl[idx]
        K_lo, K_hi = self.K_tbl[idx-1], self.K_tbl[idx]
        # print(f'v = {v}, v_lo = {v_lo}, v_hi = {v_hi}')
        # print(f'K_lo = {K_lo}, K_hi = {K_hi}')
        # 3) 线性插值
        alpha = (v - v_lo) / (v_hi - v_lo)
        return (1 - alpha) * K_lo + alpha * K_hi

    def compute_control(self, state):
        v = state.v
        K = self.get_K(v)
        # … 用 K 去算 u = -K @ (x − x_ref) 就好了
        return K


    # 计算距离自车当前位置最近的参考点
    def calc_nearest_index(self,state, cx, cy, cyaw):
        dx = [state.x - icx for icx in cx]
        dy = [state.y - icy for icy in cy]

        d = [idx ** 2 + idy ** 2 for (idx, idy) in zip(dx, dy)]

        mind = min(d)

        ind = d.index(mind)

        mind = math.sqrt(mind)

        dxl = cx[ind] - state.x
        dyl = cy[ind] - state.y
        angle = self.pi_2_pi(cyaw[ind] - math.atan2(dyl, dxl))
        if angle < 0:
            mind *= -1

        return ind, mind

    def process(self,vehicle_info,path_waypoints,speed_profile,mode,Reverse,simple,ind_list,c_path,path_re,lqr_tree,slope=0):
        cx,cy,cyaw,ck,cgear=[],[],[],[],[]
        if path_re is None:
            FLAG_RE = True
            cx = c_path[0]
            cy = c_path[1]
            cyaw = c_path[2]
            cgear = c_path[3]
            ck = c_path[4]
        else:
            FLAG_RE = False
            cx = path_re[0]
            cy = path_re[1]
            cyaw = path_re[2]
            cgear = path_re[3]
            ck = path_re[4]
        # for i in path_waypoints:
        #     cx.append(i[0])
        #     cy.append(i[1])
        #     cyaw.append(i[2])
        #     cgear.append(i[3])
        #     ck.append(i[4])
        _ego_v = vehicle_info['ego']['v_mps']
        _ego_x = vehicle_info['ego']['x']
        _ego_y = vehicle_info['ego']['y']
        _ego_yaw = vehicle_info['ego']['yaw_rad']
        state = State(_ego_x, _ego_y, _ego_yaw, _ego_v)
        # 找到最接近车辆当前位置的路标点的索引ind以及在该路标点处的横向距离e
        ind, e = self.calc_nearest_index(state, cx, cy, cyaw)
        # print(ind)
        if cgear[ind]==True:
            gear = 1
        else:
            gear = 3    # 挡位设置为3（后退）
        sp = speed_profile
        tv = sp[ind]
        k = ck[ind]
        v_state = state.v
        max_errory=0.05
        max_error_x=0.06
        key=0.2
        kex=0.35
        th_e = self.pi_2_pi(state.yaw - cyaw[ind])

        # 构建LQR表达式，X(k+1) = A * X(k) + B * u(k), 使用Riccati equation 求解LQR问题
    #     dt表示采样周期，v表示当前自车的速度
    #     A = [1.0, dt, 0.0, 0.0, 0.0
    #          0.0, 0.0, v, 0.0, 0.0]
    #          0.0, 0.0, 1.0, dt, 0.0]
    #          0.0, 0.0, 0.0, 0.0, 0.0]
    #          0.0, 0.0, 0.0, 0.0, 1.0]
        A = np.zeros((5, 5))
        A[0, 0] = 1.0
        A[0, 1] = self.dt
        A[1, 2] = v_state
        A[2, 2] = 1.0
        A[2, 3] = self.dt
        A[4, 4] = 1.0

        # 构建B矩阵，L是自车的轴距
        # B = [0.0, 0.0
        #     0.0, 0.0
        #     0.0, 0.0
        #     v/L, 0.0
        #     0.0, dt]
        B = np.zeros((5, 2))
        B[3, 0] = v_state / self.L
        B[4, 1] = self.dt

        if th_e >= 0:
            if e > 0:
                kq1 = 1
                kq3 = 0
            else:
                kq1 = 0
                kq3 = 5
        else:
            if e > 0:
                kq1 = 0
                kq3 = 5
            else:
                kq1 = 1
                kq3 = 0
        
        if abs(e)>max_errory:
           self.lqr_Q[0,0]=1
        else:
           self.lqr_Q[0,0]=1
        self.lqr_Q[1, 1] = 1

        if abs(th_e)>max_error_x:
           self.lqr_Q[2,2]=0.7
        else:
           self.lqr_Q[2,2]=kex*5**(abs(th_e))
        self.lqr_Q[3, 3] = 20
        self.lqr_Q[4, 4] =2

        # 构建了车辆动力学的状态空间模型：X(k+1) = A * X(k) + B * u(k)，X(k)是状态向量，u(k)是输入向量，A和B分别是状态转移和输入矩阵
        # 调用dlqr函数解决离散时间线性二次调节器（LQR）问题，并获得最优控制增益K
        # K, _, _ = self.dlqr(A, B, self.lqr_Q, self.lqr_R)
        K, _, _ = self.dlqr(A, B, self.lqr_Q, self.lqr_R)
        # state vector，构建状态矩阵
        # x = [e, dot_e, th_e, dot_th_e, delta_v]
        # e: lateral distance to the path， e是自车到轨迹的距离
        # dot_e: derivative of e， dot_e是自车到轨迹的距离的变化率
        # th_e: angle difference to the path， th_e是自车与期望轨迹的角度偏差
        # dot_th_e: derivative of th_e， dot_th_e是自车与期望轨迹的角度偏差的变化率
        # delta_v: difference between current speed and target speed，delta_v是当前车速与期望车速的偏差
        X = np.zeros((5, 1))
        X[0, 0] = e
        X[1, 0] = (e - self.pe) / self.dt
        X[2, 0] = th_e
        X[3, 0] = (th_e - self.pth_e) / self.dt
        X[4, 0] = v_state - tv
        # input vector，构建输入矩阵u
        # u = [delta, accel]
        # delta: steering angle，前轮转角
        # accel: acceleration，自车加速度
        ustar = -K @ X
        # 构建状态向量X并使用控制增益K计算最优输入ustar

        # 根据最优输入ustar和前馈转向角ff计算转向角delta和加速度accel。然后，该方法返回这些值以及挡位
        ff = math.atan2(self.L * k, 1)  # feedforward steering angle
        fb = self.pi_2_pi(ustar[0, 0])  # feedback steering angle
        delta = 0.01*ff + fb
        # print(delta)
        if abs(delta)>self.max_steer:
            if delta>0:
                delta = self.max_steer
            else:
                delta = -self.max_steer
        if gear == 3 and mode=="complex":
            # calc accel input
            accel = -ustar[1, 0]
        else:
            accel = ustar[1, 0]
        accel = accel + (0.02)*(-9.8)*np.sin(slope)
        if abs(accel)>15:
            if accel>0:
                accel = 15
            else:
                accel = -15

        ind_list.append(ind)
        repeats = count_repeats(ind_list)
        if repeats >= 5 and gear == 3 and FLAG_RE:
            print(111111111111111111111111111111111111)
            path_re = path_waypoints[ind + 1:]
            speed_profile = speed_profile[ind + 1:]
            cx1,cy1,cyaw1,ck1,cgear1=[],[],[],[],[]
            for i in path_re:
                cx1.append(i[0])
                cy1.append(i[1])
                cyaw1.append(i[2])
                cgear1.append(i[3])
                ck1.append(i[4])
            path_re = [cx1,cy1,cyaw1,cgear1,ck1]
        # if path_waypoints[ind][3] == 2 and repeats >= 20:
        # if (path_waypoints[ind][3] == 2 or repeats >= 50 or path_waypoints[ind][3] == 0) and simple == False :
        if path_waypoints[ind][3] == 2 and simple == False :
            # path_waypoints = path_waypoints[ind + 1:]
            # speed_profile = speed_profile[ind + 1:]
            Reverse=True
            print("进入空挡，切换控制方法为simple")
        #
        # goal_x = np.mean(observation['test_setting']["goal"]["x"])
        # goal_y = np.mean(observation['test_setting']["goal"]["y"])
        # final_dist = math.sqrt((_ego_x-goal_x) ** 2 + (_ego_y-goal_y) ** 2)
        # if repeats >= 10 and final_dist > 30:
        #     ind_list = []
        #     from planner import Planner
        #     planner = Planner(observation)
        #     print("本车超过1s未移动且距终点较远，重新规划路线")
        #     path_waypoints, speed_profile = planner.process(collision_lookup, observation,plan_all=True)
        # accel = accel + (0.02)*(-9.8)*np.sin(slope)
        print(f'ind={ind},gear={gear},reverse={Reverse},repeats={repeats}')

        return (accel,delta,gear),Reverse,ind_list,path_re,speed_profile
    